import grpc
import logging
import time

# 导入编译后的protobuf模块
import seg_match_pb2
import seg_match_pb2_grpc

# 日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class SegMatchClient:
    """gRPC客户端实现"""
    
    def __init__(self, server_address='localhost:50051'):
        """初始化客户端连接"""
        self.server_address = server_address
        self.channel = None
        self.stub = None
    
    def connect(self):
        """建立与服务器的连接"""
        logger.info(f"正在连接到gRPC服务器: {self.server_address}")
        self.channel = grpc.insecure_channel(self.server_address)
        self.stub = seg_match_pb2_grpc.SegMatchServiceStub(self.channel)
        logger.info(f"已成功连接到gRPC服务器: {self.server_address}")
    
    def close(self):
        """关闭客户端连接"""
        if self.channel:
            self.channel.close()
            logger.info(f"已关闭与gRPC服务器的连接: {self.server_address}")
    
    def bi_stream_ai(self, requests):
        """双向流调用示例"""
        if not self.stub:
            raise Exception("客户端未连接到服务器，请先调用connect()方法")
        
        logger.info("开始双向流调用BiStreamAI")
        
        try:
            # 调用双向流接口
            responses = self.stub.BiStreamAI(requests)
            
            # 处理响应流
            for response in responses:
                logger.info(f"收到响应 - status: {response.status}, message: {response.message}")
                
                # 处理分词结果
                if response.tokenize_result:
                    logger.info(f"分词结果: {response.tokenize_result.tokens}, 状态: {response.tokenize_result.status}")
                
                # 处理推荐结果
                if response.recommend_result:
                    logger.info(f"推荐结果状态: {response.recommend_result.status}")
                    for item in response.recommend_result.items:
                        logger.info(f"  推荐项: {item.name} (ID: {item.id}, 类型: {item.type}, 得分: {item.score}, 点击率: {item.click_rate})")
                        
            return True
            
        except grpc.RpcError as e:
            logger.error(f"双向流调用失败: {e.code()} - {e.details()}")
            return False
        
        except Exception as e:
            logger.error(f"双向流调用发生未知错误: {str(e)}")
            return False

def test_client():
    """测试客户端功能"""
    client = SegMatchClient('localhost:50051')
    
    try:
        # 连接服务器
        client.connect()
        
        # 准备请求数据生成器
        def request_generator():
            yield seg_match_pb2.AIRequest(
                user_input="测试分词功能",
                team_id="team_001",
                action="tokenize"
            )
            yield seg_match_pb2.AIRequest(
                user_input="注册表单",
                team_id="team_001",
                action="recommend"
            )
            yield seg_match_pb2.AIRequest(
                user_input="",
                team_id="team_001",
                action="click_rate_update",
                item_id="comp_001"
            )
        
        # 发送请求并处理响应
        client.bi_stream_ai(request_generator())
        
    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        return False
    
    finally:
        # 关闭连接
        client.close()
    
    logger.info("客户端测试完成")
    return True

if __name__ == '__main__':
    test_client()